#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Usage: result_merge.py: 合并两个 JSON Lines 格式的文件，以 base 文件为基础，用 new 文件中相同 id 的记录进行更新。
# python src/tools/result_merge.py --new "./result/results.jsonl" --base "./result/results copy.jsonl" --des "./result/results_查必准.jsonl"

import json
import argparse
from typing import Dict, List

def load_jsonl(file_path: str) -> Dict[int, dict]:
    """
    读取 JSON Lines 文件，返回以 id 为键的字典。
    """
    data = {}
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
                record_id = record.get("id")
                if record_id is not None:
                    data[record_id] = record
                else:
                    print(f"警告: 第 {line_num} 行记录缺少 'id' 字段: {line}")
            except json.JSONDecodeError as e:
                print(f"错误: 第 {line_num} 行 JSON 解析失败: {e}")
                raise
    return data

def merge_and_write(base_file: str, new_file: str, des_file: str):
    """
    以 base_file 为基础，用 new_file 中相同 id 的记录进行更新，写入 des_file。
    保持 base_file 的顺序和数量。
    """
    # 加载 new 文件数据，构建成 id -> record 的映射
    new_data = load_jsonl(new_file)
    print(f"从 '{new_file}' 加载了 {len(new_data)} 条记录。")

    processed_ids = set()
    updated_count = 0

    with open(base_file, 'r', encoding='utf-8') as base_f, \
         open(des_file, 'w', encoding='utf-8') as des_f:

        for line_num, line in enumerate(base_f, 1):
            line = line.strip()
            if not line:
                continue

            try:
                base_record = json.loads(line)
                record_id = base_record.get("id")

                if record_id is None:
                    print(f"警告: 基础文件第 {line_num} 行记录缺少 'id' 字段，使用原记录: {line}")
                    final_record = base_record
                else:
                    # 检查是否已处理过该 id（避免重复 id 导致的问题）
                    if record_id in processed_ids:
                        print(f"警告: 基础文件中 id={record_id} 重复出现，跳过后续重复项。")
                        continue
                    processed_ids.add(record_id)

                    # 如果 new 文件中存在相同 id 的记录，则使用 new 的记录
                    if record_id in new_data:
                        final_record = new_data[record_id]
                        updated_count += 1
                    else:
                        final_record = base_record

                # 写入最终记录
                des_f.write(json.dumps(final_record, ensure_ascii=False) + '\n')

            except json.JSONDecodeError as e:
                print(f"错误: 基础文件第 {line_num} 行 JSON 解析失败: {e}")
                raise

    print(f"处理完成。")
    print(f"基础文件共 {len(processed_ids)} 条有效记录。")
    print(f"更新了 {updated_count} 条记录。")
    print(f"结果已写入 '{des_file}'。")

def main():
    parser = argparse.ArgumentParser(description="合并两个 JSON Lines 文件，以 base 为基础，用 new 更新相同 id 的记录。")
    parser.add_argument('--new', required=True, help='新数据文件路径（优先使用）')
    parser.add_argument('--base', required=True, help='基础数据文件路径（保持顺序和数量）')
    parser.add_argument('--des', required=True, help='输出目标文件路径')

    args = parser.parse_args()

    merge_and_write(args.base, args.new, args.des)

if __name__ == '__main__':
    main()